import os
from sklearn.model_selection import train_test_split
from keras.callbacks import TensorBoard, ModelCheckpoint
from  util import *
np.random.seed(42)
import scipy.misc as mc
data_location = ''
#定义训练和验证数据的路径
training_images_loc = data_location + 'Crack/train/image/'
training_label_loc = data_location + 'Crack/train/label/'
validate_images_loc = data_location + 'Crack/validate/images/'
validate_label_loc = data_location + 'Crack/validate/labels/'
#获取训练和验证数据文件列表
train_files = os.listdir(training_images_loc)
train_data = []  #用于存储训练图像数据
train_label = []  #用于存储训练标签数据

validate_files = os.listdir(validate_images_loc)
validate_data = []
validate_label = []
desired_size=480

#处理训练数据
for i in train_files:
    im = mc.imread(training_images_loc + i)  #读取训练数据
    label = mc.imread(training_label_loc + "Image_" +i.split('_')[1].split(".")[0] +"_1stHO.png" )
    old_size = im.shape[:2]   #获取图像的原始尺寸
    delta_w = desired_size - old_size[1]  #计算宽度需要填充的像素
    delta_h = desired_size - old_size[0]  #计算高度需要填充的像素
    #计算上下左右需要填充的像素数
    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
    left, right = delta_w // 2, delta_w - (delta_w // 2)

    color = [0, 0, 0]  #填充颜色为黑色
    color2 = [0]  #标签图像的填充颜色
    #使用CV2库进行图像填充
    new_im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                value=color)

    new_label = cv2.copyMakeBorder(label, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                   value=color2)
    #填充后的图像调整为期望尺寸，并添加到训练数据列表
    train_data.append(cv2.resize(new_im, (desired_size, desired_size)))

    temp = cv2.resize(new_label,
                      (desired_size, desired_size))
    _, temp = cv2.threshold(temp, 127, 255, cv2.THRESH_BINARY)   #二值化处理标签图像
    train_label.append(temp)
#处理验证数据
for i in validate_files:
    im = mc.imread(validate_images_loc + i)
    label = mc.imread(validate_label_loc + "Image_" +i.split('_')[1].split(".")[0] +"_1stHO.png" )
    old_size = im.shape[:2]  # old_size is in (height, width) format
    delta_w = desired_size - old_size[1]
    delta_h = desired_size - old_size[0]

    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
    left, right = delta_w // 2, delta_w - (delta_w // 2)

    color = [0, 0, 0]
    color2 = [0]
    new_im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                value=color)

    new_label = cv2.copyMakeBorder(label, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                   value=color2)

    validate_data.append(cv2.resize(new_im, (desired_size, desired_size)))

    temp = cv2.resize(new_label,
                      (desired_size, desired_size))
    _, temp = cv2.threshold(temp, 127, 255, cv2.THRESH_BINARY)
    validate_label.append(temp)
#将训练和验证的数据转换为numpy数组
train_data = np.array(train_data)
train_label = np.array(train_label)
validate_data = np.array(validate_data)
validate_label = np.array(validate_label)

#归一化数据
x_train = train_data.astype('float32') / 255.
y_train = train_label.astype('float32') / 255.
x_validate = validate_data.astype('float32') / 255.
y_validate = validate_label.astype('float32') / 255.

#调整数据形状以适应模型的输入格式
x_train = np.reshape(x_train, (len(x_train), desired_size, desired_size, 3))  # adapt this if using `channels_first` image data format
y_train = np.reshape(y_train, (len(y_train), desired_size, desired_size, 1))  # adapt this if using `channels_first` im
x_validate = np.reshape(x_validate, (len(x_validate), desired_size, desired_size, 3))  # adapt this if using `channels_first` image data format
y_validate = np.reshape(y_validate, (len(y_validate), desired_size, desired_size, 1))  # adapt this if using `channels_first` im

#配置TensorBoard用于可视化训练过程
TensorBoard(log_dir='./autoencoder', histogram_freq=0,
            write_graph=True, write_images=True)

#导入自定义的模型
from RSAN import *
model=RSANet(input_size=(desired_size,desired_size,3),start_neurons=16,lr=1e-3,keep_prob=0.87,block_size=7)
weight="Crack/Model/RSAN.h5"   #模型权重文件路径

restore=False

if restore and os.path.isfile(weight):
    model.load_weights(weight)

#配置模型检查点的回调函数，用于训练过程保存最佳模型权重
model_checkpoint = ModelCheckpoint(weight, monitor='val_accuracy', verbose=1, save_best_only=False)

# plot_model(model, to_file='unet_resnet.png', show_shapes=False, show_layer_names=)
#训练模型
history=model.fit(x_train, y_train,
                epochs=100, #first  100 with lr=1e-3,,and last 50 with lr=1e-4
                batch_size=2,
                # validation_split=0.1,
                validation_data=(x_validate, y_validate),
                shuffle=True,
                callbacks= [TensorBoard(log_dir='./autoencoder'), model_checkpoint])
